import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import transforms,datasets
import os
from torchsummary import summary
import matplotlib.pyplot as plt

import sys
sys.path.insert(0,'F:\lab\\adv_mnist\DeepRobust-master') # 将该路径插入到包的搜索路径中
from deeprobust.image.netmodels import resnet
from deeprobust.image.attack.fgsm import FGSM
from deeprobust.image.config import attack_params

weightPath = "./trained_models/MNIST/MNIST_ResNet18_epoch_20.pt"
dataPath = './dataSet/MNIST'

get_cuda=torch.cuda.is_available()
device = torch.device("cuda" if (get_cuda and torch.cuda.is_available()) else "cpu")

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.0,), (1.0,))])
test_set = datasets.MNIST(root = './dataSet/MNIST', train=False, transform = transform, download=False)
img_data,img_label = test_set[0]
print(img_data.shape,img_label)

test_loader = torch.utils.data.DataLoader(test_set,batch_size=1,shuffle=True)
print("Test data: ",len(test_loader))

model = resnet.ResNet18([1,28,28]).to(device)

model.load_state_dict(torch.load(weightPath))
model.eval()

for batch_idx, (data, target) in enumerate(test_set):
    print(batch_idx,target)
    data, target = torch.as_tensor(data).to(device), torch.as_tensor(target).to(device)
    print(data.shape)
    F1 = FGSM(model,device='cuda')
    adv = F1.generate(data,target,**attack_params['FGSM_MNIST'])
    
    predict0 = model(data)
    predict0= predict0.argmax(dim=1, keepdim=True)
    
    predict1 = model(adv)
    predict1= predict1.argmax(dim=1, keepdim=True)
    if(batch_idx==1):
        break


